## Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
# KubeOS is licensed under the Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
## See the Mulan PSL v2 for more details.

import networkx as nx
import json
import os,time
import re,sys
import uuid
import subprocess
from argparse import ArgumentParser

def removeSOfiles(G, dotfile):
    pattern = r'\d?(?:\.\d+)?.*\.so$|so\..*\d+(?:\.\d+)?$'
    nodes = G.nodes()
    removelist = []
    for node in nodes:
        remove = True
        if node == '\\n':
            removelist.append(node)
            continue
        label = nodes[node]['label'].strip("\"")
        if re.search(pattern, label):
            for source in G.neighbors(node):
                if G.edges[source, node]['e_type'] != 'read':
                    remove = False
                    break
        else:
            remove = True

        if remove:
            removelist.append(node)
    
    G.remove_nodes_from(removelist)
    nx.drawing.nx_pydot.write_dot(G, dotfile + '.clip')
    return G

def jsonAlert(dotfile):
    G = nx.DiGraph(nx.nx_pydot.read_dot(dotfile))
    # G = removeSOfiles(G, dotfile)
    alert = {
        "alertStatus": 1,
        "alertGeneratedTime": "2024-01-01-00:00:01",
        "alertDetector": "A-SysArmor",
        "alertLevel": 1,
        "hostUUID": "000-1111",
        "hostIP": "192.168.122.234",
        "alertGeneratedTimestamp": 1234,
        "overlong": False,
        "service": "test"
    }
    alertGentime = eval(dotfile[-17:-4])
    alert["alertGeneratedTime"] = alertGentime
    alert["alertGeneratedTimestamp"] = alertGentime

    ### EDGES
    JEDGES=[]
    for u, v, data in G.edges(data=True):
        jedge = {}
        jedge["source"] = u
        jedge["target"] = v
        jedge["type"] = data['e_type']
        if 'timestamp' in data.keys():
            timestamp = data['timestamp'].split(",")[0].strip("\"")
            jedge["timestamp"] = timestamp
            jedge["host"] = data['timestamp'].split(",")[1].strip("\"")
        JEDGES.append(jedge)

    ### NODES
    JNODES=[]
    nodes = G.nodes()
    alertscore = 0
    for node in nodes:
        if node == '\\n':
            continue
        jnode = {}
        jnode["id"] = node
        nodescore = float(nodes[node]['score'].strip("\""))
        jnode["nodeMalicousScore"] = nodescore
        alertscore += nodescore
        jnode["label"] = nodes[node]['label'].strip("\"")
        shape = nodes[node]['shape'].strip("\"")
        if 'box' in shape:
            jnode["type"] = "Process"
        elif 'ellipse' in shape:
            jnode["type"] = "File"
        elif 'diamond' in shape:
            jnode["type"] = "Network"
        else:
            jnode["type"] = "unknown"
        JNODES.append(jnode)
    
    alert_graph = {}
    alert_graph["version"] = 1
    alert_graph["nodes"] = JNODES
    alert_graph["edges"] = JEDGES

    alert["maliciousScore"] =alertscore
    alert["alertGraph"] = alert_graph  ## json.dumps(alert_graph, ensure_ascii=False)
    guid = str(uuid.uuid1())
    alert["alertUUID"] = guid
    alert["nodeNum"] = len(JNODES)
    alert["edgeNum"] = len(JEDGES)
    if len(JNODES) > 200:
        alert["overlong"] = True

    newname = dotfile.rstrip('dot') + 'json'
    with open(newname, 'w') as f:
        json.dump(alert, f, indent=4)
    return G

if __name__ == "__main__":

    parser = ArgumentParser(description="DOT alert to JSON alert")
    parser.add_argument("--falert", type=str, default='alert', help="alert DOT file path")
    args = parser.parse_args()

    if os.path.exists(args.falert):
        jsonAlert(args.falert)
    else:
        print("alert DOT file not exist, check alert path")